package main

import (
	"context"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"
)

// SOCKS5协议常量
const (
	Socks5Version = 0x05
	
	// 认证方法
	Socks5AuthNone     = 0x00
	Socks5AuthPassword = 0x02
	Socks5AuthNoAcceptable = 0xFF
	
	// 命令类型
	Socks5CmdConnect      = 0x01
	Socks5CmdBind         = 0x02
	Socks5CmdUDPAssociate = 0x03
	
	// 地址类型
	Socks5AddrTypeIPv4   = 0x01
	Socks5AddrTypeDomain = 0x03
	Socks5AddrTypeIPv6   = 0x04
	
	// 响应状态
	Socks5StatusSuccess = 0x00
	Socks5StatusFailure = 0x01
	Socks5StatusNotAllowed = 0x02
	Socks5StatusNetworkUnreachable = 0x03
	Socks5StatusHostUnreachable = 0x04
	Socks5StatusConnectionRefused = 0x05
	Socks5StatusTTLExpired = 0x06
	Socks5StatusCommandNotSupported = 0x07
	Socks5StatusAddressTypeNotSupported = 0x08
)

// 定义命令行参数
var (
	proxyHost     string
	proxyPort     int
	proxyUser     string
	proxyPassword string
	listenPort    int
	proxyType     string
	socks5Listen  bool
	help          bool
	timeoutMinutes int
)

// 初始化命令行参数
func init() {
	flag.StringVar(&proxyHost, "host", "", "代理服务器主机地址")
	flag.IntVar(&proxyPort, "port", 0, "代理服务器端口")
	flag.StringVar(&proxyUser, "user", "", "代理服务器用户名（可选）")
	flag.StringVar(&proxyPassword, "password", "", "代理服务器密码（可选）")
	flag.IntVar(&listenPort, "listen", 8080, "本地监听端口")
	flag.StringVar(&proxyType, "type", "http", "代理类型: http 或 socks5")
	flag.BoolVar(&socks5Listen, "socks5-server", false, "启动SOCKS5服务器模式（不转发到上游代理）")
	flag.IntVar(&timeoutMinutes, "timeout", 3, "程序运行的最大时长（分钟），超时后自动关闭")
	flag.BoolVar(&help, "help", false, "显示帮助信息")
	flag.BoolVar(&help, "h", false, "显示帮助信息（简写）")
}

// 显示使用说明
func showUsage() {
	fmt.Println("用法: proxy -host <代理主机> -port <代理端口> [选项...]")
	fmt.Println("\n选项:")
	flag.PrintDefaults()
	fmt.Println("\n示例:")
	fmt.Println("  HTTP代理:   proxy -host 10.0.0.1 -port 8888 -type http -listen 8080")
	fmt.Println("  SOCKS5代理: proxy -host 10.0.0.1 -port 8888 -type socks5 -listen 8080")
	fmt.Println("  带认证代理: proxy -host 10.0.0.1 -port 8888 -user username -password pwd123 -listen 8080")
	fmt.Println("  SOCKS5服务器: proxy -socks5-server -listen 8080")
	fmt.Println("\n说明: 此程序将在本地启动一个代理服务器，将所有请求转发到指定的上游代理服务器")
}

func main() {
	// 解析命令行参数
	flag.Parse()

	// 显示帮助信息
	if help || (!socks5Listen && (proxyHost == "" || proxyPort == 0)) {
		showUsage()
		os.Exit(0)
	}
	
	// 验证代理类型
	if proxyType != "http" && proxyType != "socks5" {
		log.Fatalf("不支持的代理类型: %s, 只支持 http 或 socks5", proxyType)
	}

	// 设置程序超时自动关闭
	if timeoutMinutes > 0 {
		log.Printf("程序将在 %d 分钟后自动关闭", timeoutMinutes)
		go func() {
			timeout := time.Duration(timeoutMinutes) * time.Minute
			timer := time.NewTimer(timeout)
			<-timer.C
			log.Println("程序运行时间已达到设定值，正在自动关闭...")
			// 在Windows上直接调用os.Exit()来关闭程序
			os.Exit(0)
		}()
	}

	// 如果是SOCKS5服务器模式，直接启动SOCKS5服务器
	if socks5Listen {
		startSocks5Server(listenPort)
		return
	}
	
	// 构建代理URL
	var proxyUrlStr string
	if proxyType == "http" {
		proxyUrlStr = fmt.Sprintf("http://%s:%d", proxyHost, proxyPort)
		if proxyUser != "" && proxyPassword != "" {
			proxyUrlStr = fmt.Sprintf("http://%s:%s@%s:%d", proxyUser, proxyPassword, proxyHost, proxyPort)
		}
	} else { // socks5
		proxyUrlStr = fmt.Sprintf("socks5://%s:%d", proxyHost, proxyPort)
		if proxyUser != "" && proxyPassword != "" {
			proxyUrlStr = fmt.Sprintf("socks5://%s:%s@%s:%d", proxyUser, proxyPassword, proxyHost, proxyPort)
		}
	}

	// 解析代理URL
	proxyUrl, err := url.Parse(proxyUrlStr)
	if err != nil {
		log.Fatalf("无效的代理URL: %v", err)
	}

	// 创建自定义Transport
	transport := &http.Transport{
		Proxy: http.ProxyURL(proxyUrl), // 设置代理认证
		DialContext: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).DialContext,
		MaxIdleConns:          100,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
		ResponseHeaderTimeout: 30 * time.Second,
	}

	// 配置反向代理
	reverseProxy := &httputil.ReverseProxy{
		Transport: transport,
		Director: func(req *http.Request) {
			// 保持原始请求信息
			req.URL.Scheme = "http"
			req.URL.Host = req.Host
			// 添加必要头信息
			req.Header.Set("X-Forwarded-Host", req.Host)
		},
	}

	// 创建HTTP服务器
	listenAddr := fmt.Sprintf(":%d", listenPort)
	server := &http.Server{
		Addr:    listenAddr, // 中转服务监听端口
		Handler: http.HandlerFunc(handler(reverseProxy)),
	}

	log.Printf("启动代理服务器在 %s，转发到 %s", listenAddr, proxyUrlStr)
	
	// 优雅关闭服务器
	go func() {
		sigChan := make(chan os.Signal, 1)
		signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
		<-sigChan
		log.Println("接收到关闭信号，正在关闭服务器...")
		
		// 创建一个5秒的上下文用于关闭服务器
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		
		if err := server.Shutdown(ctx); err != nil {
			log.Fatalf("服务器关闭错误: %v", err)
		}
	}()
	
	// 启动服务器
	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		log.Fatalf("服务器启动错误: %v", err)
	}
}

// 处理所有请求
func handler(p *httputil.ReverseProxy) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		// 记录请求信息
		log.Printf("收到请求: %s %s %s", r.Method, r.URL.String(), r.Proto)

		// 普通HTTP请求
		if r.Method != "CONNECT" {
			p.ServeHTTP(w, r)
			return
		}

		// 处理HTTPS/WebSocket等CONNECT请求
		serveConnect(w, r, p.Transport)
	}
}

// 处理CONNECT隧道
func serveConnect(w http.ResponseWriter, r *http.Request, transport http.RoundTripper) {
	// 获取目标地址
	target := r.URL.Host
	if target == "" {
		http.Error(w, "无效的目标地址", http.StatusBadRequest)
		log.Printf("CONNECT请求缺少目标地址: %s", r.URL)
		return
	}

	// 根据代理类型选择不同的连接方式
	var destConn net.Conn
	var err error
	
	if proxyType == "socks5" {
		// 使用SOCKS5代理连接
		destConn, err = dialViaSocks5(target)
	} else if proxyType == "http" && proxyHost != "" && proxyPort > 0 {
		// 使用HTTP代理连接
		// 构建代理URL
		proxyUrlStr := fmt.Sprintf("http://%s:%d", proxyHost, proxyPort)
		if proxyUser != "" && proxyPassword != "" {
			proxyUrlStr = fmt.Sprintf("http://%s:%s@%s:%d", proxyUser, proxyPassword, proxyHost, proxyPort)
		}
		
		// 连接到HTTP代理
		proxyURL, err := url.Parse(proxyUrlStr)
		if err != nil {
			http.Error(w, "代理URL解析错误", http.StatusInternalServerError)
			log.Printf("解析代理URL失败: %v", err)
			return
		}
		
		// 创建到代理的连接
		proxyConn, err := net.DialTimeout("tcp", proxyURL.Host, 30*time.Second)
		if err != nil {
			http.Error(w, "连接代理失败", http.StatusServiceUnavailable)
			log.Printf("连接到HTTP代理失败: %v", err)
			return
		}
		
		// 发送CONNECT请求到代理
		connectReq := fmt.Sprintf(
			"CONNECT %s HTTP/1.1\r\n"+
			"Host: %s\r\n", 
			target, target)
		
		// 添加代理认证
		if proxyUser != "" && proxyPassword != "" {
			auth := base64.StdEncoding.EncodeToString([]byte(proxyUser + ":" + proxyPassword))
			connectReq += fmt.Sprintf("Proxy-Authorization: Basic %s\r\n", auth)
		}
		
		// 结束请求头
		connectReq += "\r\n"
		
		// 发送请求
		if _, err = proxyConn.Write([]byte(connectReq)); err != nil {
			proxyConn.Close()
			http.Error(w, "发送CONNECT请求失败", http.StatusServiceUnavailable)
			log.Printf("发送CONNECT请求到代理失败: %v", err)
			return
		}
		
		// 读取代理响应
		buf := make([]byte, 4096)
		deadline := time.Now().Add(10 * time.Second)
		proxyConn.SetReadDeadline(deadline)
		
		n, err := proxyConn.Read(buf)
		if err != nil {
			proxyConn.Close()
			http.Error(w, "读取代理响应失败", http.StatusServiceUnavailable)
			log.Printf("读取代理响应失败: %v", err)
			return
		}
		
		// 重置读取超时
		proxyConn.SetReadDeadline(time.Time{})
		
		// 解析响应
		response := string(buf[:n])
		if !strings.Contains(response, "200") {
			proxyConn.Close()
			http.Error(w, "代理连接失败", http.StatusServiceUnavailable)
			log.Printf("代理返回非200响应: %s", response)
			return
		}
		
		// 使用代理连接作为目标连接
		destConn = proxyConn
	} else {
		// 直接连接到目标
		transportPtr, ok := transport.(*http.Transport)
		if !ok {
			http.Error(w, "内部服务器错误", http.StatusInternalServerError)
			log.Printf("无法转换transport类型")
			return
		}
		
		if transportPtr != nil && transportPtr.Dial != nil {
			destConn, err = transportPtr.Dial("tcp", target)
		} else {
			// 如果Dial函数未设置，使用默认的拨号器
			destConn, err = net.DialTimeout("tcp", target, 30*time.Second)
		}
	}
	
	if err != nil {
		http.Error(w, "连接失败", http.StatusServiceUnavailable)
		log.Printf("无法连接到目标: %s, 错误: %v", target, err)
		return
	}

	// 响应客户端连接就绪
	w.WriteHeader(http.StatusOK)

	// 劫持连接进行全双工通信
	hijacker, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "Hijacking not supported", http.StatusInternalServerError)
		return
	}

	clientConn, _, err := hijacker.Hijack()
	if err != nil {
		http.Error(w, "Hijack failed", http.StatusServiceUnavailable)
		return
	}

	// 双向数据转发
	go transfer(destConn, clientConn)
	go transfer(clientConn, destConn)
}

// 双向数据转发
func transfer(dst, src net.Conn) {
	defer func() {
		if err := dst.Close(); err != nil {
			log.Printf("关闭连接错误: %v", err)
		}
	}()
	defer func() {
		if err := src.Close(); err != nil {
			log.Printf("关闭连接错误: %v", err)
		}
	}()
	
	// 添加错误处理
	buffer := make([]byte, 32*1024) // 32KB 缓冲区
	n, err := io.CopyBuffer(dst, src, buffer)
	if err != nil && err != io.EOF {
		log.Printf("数据传输错误: %v, 已传输 %d 字节", err, n)
	}
}

// 启动SOCKS5服务器
func startSocks5Server(port int) {
	addr := fmt.Sprintf(":%d", port)
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatalf("无法监听端口 %s: %v", addr, err)
	}
	defer listener.Close()
	
	log.Printf("SOCKS5服务器启动在 %s", addr)
	
	// 设置优雅关闭
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	
	go func() {
		<-sigChan
		log.Println("接收到关闭信号，正在关闭SOCKS5服务器...")
		listener.Close()
	}()
	
	for {
		conn, err := listener.Accept()
		if err != nil {
			if errors.Is(err, net.ErrClosed) {
				log.Println("服务器已关闭")
				break
			}
			log.Printf("接受连接错误: %v", err)
			continue
		}
		
		go handleSocks5Client(conn)
	}
}

// 处理SOCKS5客户端连接
func handleSocks5Client(client net.Conn) {
	defer func() {
		if err := client.Close(); err != nil {
			log.Printf("关闭客户端连接错误: %v", err)
		}
	}()
	
	// 设置超时
	if err := client.SetDeadline(time.Now().Add(30 * time.Second)); err != nil {
		log.Printf("设置超时错误: %v", err)
		return
	}
	
	// 1. 协商认证方法
	if err := negotiateAuth(client); err != nil {
		log.Printf("认证协商失败: %v", err)
		return
	}
	
	// 2. 处理连接请求
	target, err := handleRequest(client)
	if err != nil {
		log.Printf("处理请求失败: %v", err)
		return
	}
	
	// 3. 转发数据
	go transfer(target, client)
	go transfer(client, target)
	
	// 清除超时设置
	if err := client.SetDeadline(time.Time{}); err != nil {
		log.Printf("清除超时设置错误: %v", err)
	}
}

// 协商认证方法
func negotiateAuth(conn net.Conn) error {
	// 读取客户端支持的认证方法
	buf := make([]byte, 258) // 1字节版本号 + 1字节方法数量 + 最多256种方法
	
	// 读取版本和方法数量
	if _, err := io.ReadFull(conn, buf[:2]); err != nil {
		return fmt.Errorf("读取认证头失败: %v", err)
	}
	
	ver, nmethods := buf[0], buf[1]
	if ver != Socks5Version {
		return fmt.Errorf("不支持的SOCKS版本: %d", ver)
	}
	
	// 读取方法列表
	if _, err := io.ReadFull(conn, buf[2:2+nmethods]); err != nil {
		return fmt.Errorf("读取认证方法失败: %v", err)
	}
	
	// 检查是否支持无认证方式或用户名密码认证
	hasNoAuth := false
	hasPasswordAuth := false
	for i := 0; i < int(nmethods); i++ {
		if buf[2+i] == Socks5AuthNone {
			hasNoAuth = true
		}
		if buf[2+i] == Socks5AuthPassword {
			hasPasswordAuth = true
		}
	}
	
	// 如果有用户名和密码，优先使用密码认证
	if proxyUser != "" && proxyPassword != "" && hasPasswordAuth {
		// 选择密码认证方式
		if _, err := conn.Write([]byte{Socks5Version, Socks5AuthPassword}); err != nil {
			return fmt.Errorf("发送认证响应失败: %v", err)
		}
		
		// 执行用户名密码认证
		return performPasswordAuth(conn, proxyUser, proxyPassword)
	}
	
	// 响应客户端
	if hasNoAuth {
		// 选择无认证方式
		if _, err := conn.Write([]byte{Socks5Version, Socks5AuthNone}); err != nil {
			return fmt.Errorf("发送认证响应失败: %v", err)
		}
		return nil
	}
	
	// 不支持客户端提供的任何认证方法
	if _, err := conn.Write([]byte{Socks5Version, Socks5AuthNoAcceptable}); err != nil {
		return fmt.Errorf("发送认证拒绝失败: %v", err)
	}
	return errors.New("不支持客户端提供的认证方法")
}

// 执行用户名密码认证
func performPasswordAuth(conn net.Conn, username, password string) error {
	// 读取认证请求
	buf := make([]byte, 513) // 1字节版本 + 1字节用户名长度 + 最多255字节用户名 + 1字节密码长度 + 最多255字节密码
	
	// 读取版本和用户名长度
	if _, err := io.ReadFull(conn, buf[:2]); err != nil {
		return fmt.Errorf("读取认证版本和用户名长度失败: %v", err)
	}
	
	ver, userLen := buf[0], buf[1]
	if ver != 0x01 { // 用户名密码认证的版本号是0x01
		return fmt.Errorf("不支持的认证版本: %d", ver)
	}
	
	// 读取用户名
	if _, err := io.ReadFull(conn, buf[2:2+userLen]); err != nil {
		return fmt.Errorf("读取用户名失败: %v", err)
	}
	clientUsername := string(buf[2:2+userLen])
	
	// 读取密码长度
	if _, err := io.ReadFull(conn, buf[2+userLen:2+userLen+1]); err != nil {
		return fmt.Errorf("读取密码长度失败: %v", err)
	}
	passLen := buf[2+userLen]
	
	// 读取密码
	if _, err := io.ReadFull(conn, buf[2+userLen+1:2+userLen+1+passLen]); err != nil {
		return fmt.Errorf("读取密码失败: %v", err)
	}
	clientPassword := string(buf[2+userLen+1:2+userLen+1+passLen])
	
	// 验证用户名和密码
	if clientUsername == username && clientPassword == password {
		// 认证成功
		if _, err := conn.Write([]byte{0x01, 0x00}); err != nil {
			return fmt.Errorf("发送认证成功响应失败: %v", err)
		}
		return nil
	}
	
	// 认证失败
	if _, err := conn.Write([]byte{0x01, 0x01}); err != nil {
		return fmt.Errorf("发送认证失败响应失败: %v", err)
	}
	return errors.New("用户名或密码不正确")
}

// 通过SOCKS5代理连接到目标
func dialViaSocks5(target string) (net.Conn, error) {
	// 解析目标地址
	host, portStr, err := net.SplitHostPort(target)
	if err != nil {
		return nil, fmt.Errorf("解析目标地址失败: %v", err)
	}
	
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, fmt.Errorf("解析端口失败: %v", err)
	}
	
	// 连接到SOCKS5代理服务器
	proxyAddr := fmt.Sprintf("%s:%d", proxyHost, proxyPort)
	conn, err := net.DialTimeout("tcp", proxyAddr, 10*time.Second)
	if err != nil {
		return nil, fmt.Errorf("连接到SOCKS5代理服务器失败: %v", err)
	}
	
	// 设置超时
	if err := conn.SetDeadline(time.Now().Add(30 * time.Second)); err != nil {
		conn.Close()
		return nil, fmt.Errorf("设置超时失败: %v", err)
	}
	
	// 1. 协商认证方法
	if err := socks5ClientAuth(conn); err != nil {
		conn.Close()
		return nil, fmt.Errorf("SOCKS5认证失败: %v", err)
	}
	
	// 2. 发送连接请求
	if err := socks5ClientConnect(conn, host, port); err != nil {
		conn.Close()
		return nil, fmt.Errorf("SOCKS5连接请求失败: %v", err)
	}
	
	// 清除超时设置
	if err := conn.SetDeadline(time.Time{}); err != nil {
		conn.Close()
		return nil, fmt.Errorf("清除超时设置失败: %v", err)
	}
	
	return conn, nil
}

// SOCKS5客户端认证
func socks5ClientAuth(conn net.Conn) error {
	// 构建认证请求
	var request []byte
	
	if proxyUser != "" && proxyPassword != "" {
		// 支持无认证和用户名密码认证
		request = []byte{Socks5Version, 2, Socks5AuthNone, Socks5AuthPassword}
	} else {
		// 只支持无认证
		request = []byte{Socks5Version, 1, Socks5AuthNone}
	}
	
	// 发送认证请求
	if _, err := conn.Write(request); err != nil {
		return fmt.Errorf("发送认证请求失败: %v", err)
	}
	
	// 读取服务器响应
	response := make([]byte, 2)
	if _, err := io.ReadFull(conn, response); err != nil {
		return fmt.Errorf("读取认证响应失败: %v", err)
	}
	
	// 检查版本和认证方法
	if response[0] != Socks5Version {
		return fmt.Errorf("不支持的SOCKS版本: %d", response[0])
	}
	
	// 处理不同的认证方法
	switch response[1] {
	case Socks5AuthNone:
		// 无需认证
		return nil
		
	case Socks5AuthPassword:
		// 需要用户名密码认证
		if proxyUser == "" || proxyPassword == "" {
			return errors.New("服务器要求用户名密码认证，但未提供凭据")
		}
		return socks5ClientPasswordAuth(conn, proxyUser, proxyPassword)
		
	case Socks5AuthNoAcceptable:
		return errors.New("服务器不接受任何提供的认证方法")
		
	default:
		return fmt.Errorf("不支持的认证方法: %d", response[1])
	}
}

// SOCKS5客户端用户名密码认证
func socks5ClientPasswordAuth(conn net.Conn, username, password string) error {
	// 检查用户名和密码长度
	if len(username) > 255 || len(password) > 255 {
		return errors.New("用户名或密码过长")
	}
	
	// 构建认证请求
	request := make([]byte, 0, 3+len(username)+len(password))
	request = append(request, 0x01) // 认证协议版本
	request = append(request, byte(len(username)))
	request = append(request, []byte(username)...)
	request = append(request, byte(len(password)))
	request = append(request, []byte(password)...)
	
	// 发送认证请求
	if _, err := conn.Write(request); err != nil {
		return fmt.Errorf("发送密码认证请求失败: %v", err)
	}
	
	// 读取服务器响应
	response := make([]byte, 2)
	if _, err := io.ReadFull(conn, response); err != nil {
		return fmt.Errorf("读取密码认证响应失败: %v", err)
	}
	
	// 检查认证结果
	if response[0] != 0x01 {
		return fmt.Errorf("不支持的认证协议版本: %d", response[0])
	}
	
	if response[1] != 0x00 {
		return errors.New("用户名密码认证失败")
	}
	
	return nil
}

// SOCKS5客户端连接请求
func socks5ClientConnect(conn net.Conn, host string, port int) error {
	// 构建连接请求
	request := make([]byte, 0, 10+len(host))
	request = append(request, Socks5Version, Socks5CmdConnect, 0x00) // 版本、命令、保留字段
	
	// 添加目标地址
	ip := net.ParseIP(host)
	if ip == nil {
		// 域名地址
		if len(host) > 255 {
			return errors.New("域名过长")
		}
		request = append(request, Socks5AddrTypeDomain, byte(len(host)))
		request = append(request, []byte(host)...)
	} else if ip4 := ip.To4(); ip4 != nil {
		// IPv4地址
		request = append(request, Socks5AddrTypeIPv4)
		request = append(request, ip4...)
	} else if ip6 := ip.To16(); ip6 != nil {
		// IPv6地址
		request = append(request, Socks5AddrTypeIPv6)
		request = append(request, ip6...)
	} else {
		return errors.New("无效的IP地址")
	}
	
	// 添加端口（网络字节序）
	portBytes := make([]byte, 2)
	binary.BigEndian.PutUint16(portBytes, uint16(port))
	request = append(request, portBytes...)
	
	// 发送连接请求
	if _, err := conn.Write(request); err != nil {
		return fmt.Errorf("发送连接请求失败: %v", err)
	}
	
	// 读取服务器响应
	response := make([]byte, 4)
	if _, err := io.ReadFull(conn, response); err != nil {
		return fmt.Errorf("读取连接响应失败: %v", err)
	}
	
	// 检查版本和状态
	if response[0] != Socks5Version {
		return fmt.Errorf("不支持的SOCKS版本: %d", response[0])
	}
	
	if response[1] != Socks5StatusSuccess {
		var errMsg string
		switch response[1] {
		case Socks5StatusFailure:
			errMsg = "一般性失败"
		case Socks5StatusNotAllowed:
			errMsg = "规则集不允许连接"
		case Socks5StatusNetworkUnreachable:
			errMsg = "网络不可达"
		case Socks5StatusHostUnreachable:
			errMsg = "主机不可达"
		case Socks5StatusConnectionRefused:
			errMsg = "连接被拒绝"
		case Socks5StatusTTLExpired:
			errMsg = "TTL已过期"
		case Socks5StatusCommandNotSupported:
			errMsg = "不支持的命令"
		case Socks5StatusAddressTypeNotSupported:
			errMsg = "不支持的地址类型"
		default:
			errMsg = fmt.Sprintf("未知错误: %d", response[1])
		}
		return fmt.Errorf("SOCKS5连接失败: %s", errMsg)
	}
	
	// 读取绑定地址类型
	atyp := response[3]
	
	// 跳过绑定地址和端口（我们不需要使用它们）
	var skipBytes int
	switch atyp {
	case Socks5AddrTypeIPv4:
		skipBytes = 4 + 2 // IPv4(4) + 端口(2)
	case Socks5AddrTypeDomain:
		// 读取域名长度
		domainLen := make([]byte, 1)
		if _, err := io.ReadFull(conn, domainLen); err != nil {
			return fmt.Errorf("读取域名长度失败: %v", err)
		}
		skipBytes = int(domainLen[0]) + 2 // 域名 + 端口(2)
	case Socks5AddrTypeIPv6:
		skipBytes = 16 + 2 // IPv6(16) + 端口(2)
	default:
		return fmt.Errorf("不支持的地址类型: %d", atyp)
	}
	
	// 跳过剩余字节
	if skipBytes > 0 {
		skip := make([]byte, skipBytes)
		if _, err := io.ReadFull(conn, skip); err != nil {
			return fmt.Errorf("读取绑定地址和端口失败: %v", err)
		}
	}
	
	return nil
}

// 处理连接请求
func handleRequest(conn net.Conn) (net.Conn, error) {
	buf := make([]byte, 4) // 版本(1) + 命令(1) + 保留字段(1) + 地址类型(1)
	
	// 读取请求头
	if _, err := io.ReadFull(conn, buf); err != nil {
		return nil, fmt.Errorf("读取请求头失败: %v", err)
	}
	
	ver, cmd, _, atyp := buf[0], buf[1], buf[2], buf[3]
	if ver != Socks5Version {
		return nil, fmt.Errorf("不支持的SOCKS版本: %d", ver)
	}
	
	// 目前只支持CONNECT命令
	if cmd != Socks5CmdConnect {
		sendReply(conn, Socks5StatusCommandNotSupported, nil)
		return nil, fmt.Errorf("不支持的命令: %d", cmd)
	}
	
	// 解析目标地址
	var host string
	var err error
	
	switch atyp {
	case Socks5AddrTypeIPv4:
		buf = make([]byte, 4)
		if _, err := io.ReadFull(conn, buf); err != nil {
			return nil, fmt.Errorf("读取IPv4地址失败: %v", err)
		}
		host = net.IP(buf).String()
		
	case Socks5AddrTypeDomain:
		buf = make([]byte, 1)
		if _, err := io.ReadFull(conn, buf); err != nil {
			return nil, fmt.Errorf("读取域名长度失败: %v", err)
		}
		domainLen := buf[0]
		
		buf = make([]byte, domainLen)
		if _, err := io.ReadFull(conn, buf); err != nil {
			return nil, fmt.Errorf("读取域名失败: %v", err)
		}
		host = string(buf)
		
	case Socks5AddrTypeIPv6:
		buf = make([]byte, 16)
		if _, err := io.ReadFull(conn, buf); err != nil {
			return nil, fmt.Errorf("读取IPv6地址失败: %v", err)
		}
		host = net.IP(buf).String()
		
	default:
		sendReply(conn, Socks5StatusAddressTypeNotSupported, nil)
		return nil, fmt.Errorf("不支持的地址类型: %d", atyp)
	}
	
	// 读取端口
	buf = make([]byte, 2)
	if _, err := io.ReadFull(conn, buf); err != nil {
		return nil, fmt.Errorf("读取端口失败: %v", err)
	}
	port := binary.BigEndian.Uint16(buf)
	
	// 目标地址
	target := net.JoinHostPort(host, strconv.Itoa(int(port)))
	log.Printf("SOCKS5连接请求: %s", target)
	
	// 连接到目标服务器
	remoteConn, err := net.DialTimeout("tcp", target, 10*time.Second)
	if err != nil {
		// 根据错误类型返回不同的状态码
		var status byte = Socks5StatusFailure
		if strings.Contains(err.Error(), "connection refused") {
			status = Socks5StatusConnectionRefused
		} else if strings.Contains(err.Error(), "network is unreachable") {
			status = Socks5StatusNetworkUnreachable
		} else if strings.Contains(err.Error(), "no route to host") {
			status = Socks5StatusHostUnreachable
		}
		
		sendReply(conn, status, nil)
		return nil, fmt.Errorf("连接目标服务器失败: %v", err)
	}
	
	// 获取本地地址信息
	localAddr := remoteConn.LocalAddr().(*net.TCPAddr)
	
	// 发送成功响应
	if err := sendReply(conn, Socks5StatusSuccess, localAddr); err != nil {
		remoteConn.Close()
		return nil, fmt.Errorf("发送响应失败: %v", err)
	}
	
	return remoteConn, nil
}

// 发送SOCKS5响应
func sendReply(conn net.Conn, status byte, addr *net.TCPAddr) error {
	// 如果没有提供地址，使用空地址
	if addr == nil {
		addr = &net.TCPAddr{
			IP:   net.IPv4zero,
			Port: 0,
		}
	}
	
	// 构建响应
	response := make([]byte, 0, 10)
	response = append(response, Socks5Version, status, 0x00 /* 保留字段 */)
	
	// 添加绑定地址
	if ip4 := addr.IP.To4(); ip4 != nil {
		response = append(response, Socks5AddrTypeIPv4)
		response = append(response, ip4...)
	} else if ip6 := addr.IP.To16(); ip6 != nil {
		response = append(response, Socks5AddrTypeIPv6)
		response = append(response, ip6...)
	} else {
		return errors.New("无效的IP地址")
	}
	
	// 添加端口（网络字节序）
	portBytes := make([]byte, 2)
	binary.BigEndian.PutUint16(portBytes, uint16(addr.Port))
	response = append(response, portBytes...)
	
	// 发送响应
	_, err := conn.Write(response)
	return err
}
